{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mUCuTnh1vDFA"
      },
      "source": [
        "Copyright 2022 Google LLC\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "you may not use this file except in compliance with the License.\n",
        "You may obtain a copy of the License at\n",
        "\n",
        "    https://www.apache.org/licenses/LICENSE-2.0\n",
        "\n",
        "Unless required by applicable law or agreed to in writing, software\n",
        "distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "See the License for the specific language governing permissions and\n",
        "limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NPxG768ivKvD"
      },
      "source": [
        "## Discrete simulation study\n",
        "\n",
        "An implementation of the simulation study for discrete observations described in https://arxiv.org/abs/2212.11254. This notebook implements Algorithm 1 described in the paper to produce the results included in Table 1.\n",
        "\n",
        "This notebook relies on previously executing `colab/synthetic_data_to_file.ipynb`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gsxLNGaIW1YY"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "import pandas as pd\n",
        "import matplotlib.pyplot as plt\n",
        "import warnings\n",
        "import re\n",
        "import os\n",
        "from sklearn.cluster import KMeans\n",
        "\n",
        "with warnings.catch_warnings():\n",
        "  warnings.filterwarnings('ignore', category=DeprecationWarning)\n",
        "  import scipy.spatial"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2-GcsPRhXxbF"
      },
      "outputs": [],
      "source": [
        "def rmse(pred, true):\n",
        "  return np.sqrt(np.mean((pred - true)**2))\n",
        "\n",
        "\n",
        "def extract_from_df_nested(samples_df, cols=['u', 'x', 'w', 'c', 'c_logits', 'y', 'y_logits', 'y_one_hot', 'w_binary', 'w_one_hot', 'u_one_hot', 'x_scaled']):\n",
        "  \"\"\"\n",
        "  Extracts nested dict of numpy arrays from dataframe with structure {domain: {partition: data}}\n",
        "  \"\"\"\n",
        "  result = {}\n",
        "  if 'domain' in samples_df.keys():\n",
        "    for domain in samples_df['domain'].unique():\n",
        "      result[domain] = {}\n",
        "      domain_df = samples_df.query('domain == @domain')\n",
        "      for partition in domain_df['partition'].unique():\n",
        "        partition_df = domain_df.query('partition == @partition')\n",
        "        result[domain][partition] = extract_from_df(partition_df, cols=cols)\n",
        "  else:\n",
        "    for partition in samples_df['partition'].unique():\n",
        "        partition_df = samples_df.query('partition == @partition')\n",
        "        result[partition] = extract_from_df(partition_df, cols=cols)\n",
        "  return result\n",
        "\n",
        "\n",
        "def extract_from_df(samples_df,\n",
        "                    cols=[\n",
        "                        'u', 'x', 'w', 'c', 'c_logits', 'y', 'y_logits',\n",
        "                        'y_one_hot', 'w_binary', 'w_one_hot', 'u_one_hot',\n",
        "                        'x_scaled'\n",
        "                    ]):\n",
        "  \"\"\"Extracts dict of numpy arrays from dataframe\"\"\"\n",
        "  result = {}\n",
        "  for col in cols:\n",
        "    if col in samples_df.columns:\n",
        "      result[col] = samples_df[col].values\n",
        "    else:\n",
        "      match_str = f'^{col}_\\d$'\n",
        "      r = re.compile(match_str, re.IGNORECASE)\n",
        "      matching_columns = list(filter(r.match, samples_df.columns))\n",
        "      if len(matching_columns) == 0:\n",
        "        continue\n",
        "      result[col] = samples_df[matching_columns].to_numpy()\n",
        "  return result\n",
        "\n",
        "\n",
        "# load data and discretize it\n",
        "def discretize_data(data_dict, is_univariate=True):\n",
        "\n",
        "  X = np.concatenate(\n",
        "      (data_dict['source']['train']['x'], data_dict['source']['val']['x']))\n",
        "  Y = np.concatenate(\n",
        "      (data_dict['source']['train']['y'], data_dict['source']['val']['y']))\n",
        "  X_shift = np.concatenate(\n",
        "      (data_dict['target']['train']['x'], data_dict['target']['val']['x']))\n",
        "  Y_shift = np.concatenate(\n",
        "      (data_dict['target']['train']['y'], data_dict['target']['val']['y']))\n",
        "  X_shift_test = data_dict['target']['test']['x']\n",
        "  Y_shift_test = data_dict['target']['test']['y']\n",
        "  C = np.concatenate(\n",
        "      (data_dict['source']['train']['c'], data_dict['source']['val']['c']))\n",
        "  if is_univariate:\n",
        "    quant_X = np.quantile(X, [0.25, 0.5, 0.75])\n",
        "    X = np.digitize(X, quant_X)\n",
        "    X_shift = np.digitize(X_shift, quant_X)\n",
        "    X_shift_test = np.digitize(X_shift_test, quant_X)\n",
        "  else:\n",
        "    kmeans = KMeans(n_clusters=2, random_state=0).fit(X)\n",
        "    X = kmeans.labels_\n",
        "    X_shift = kmeans.predict(X_shift)\n",
        "    X_shift_test = kmeans.predict(X_shift_test)\n",
        "    C = np.dot(C, np.array([4, 2, 1]))  # treat multi-hot C as a binary encoding\n",
        "  W = np.concatenate((data_dict['source']['train']['w_binary'], data_dict['source']['val']['w_binary']))\n",
        "  U = np.concatenate(\n",
        "      (data_dict['source']['train']['u'], data_dict['source']['val']['u']))\n",
        "  return X, C, Y, W, U, X_shift, Y_shift, X_shift_test, Y_shift_test\n",
        "\n",
        "\n",
        "def estimate_q_y_x(X, C, Y, W, U, X_shift, Y_shift, X_shift_test, Y_shift_test):\n",
        "  k_W = len(np.unique(W))\n",
        "  k_U = len(np.unique(U))\n",
        "  k_X = len(np.unique(X))\n",
        "  k_C = len(np.unique(C))\n",
        "  k_Y = len(np.unique(Y))\n",
        "  small = 1e-9\n",
        "\n",
        "  # need to fix a c and y index, just use most popular ones\n",
        "  c_index = 0\n",
        "  c_sum = 0\n",
        "  for c in np.unique(C):\n",
        "    if np.sum(C == c) \u003e c_sum:\n",
        "      c_sum = np.sum(C == c)\n",
        "      c_index = c\n",
        "\n",
        "  y_index = 0\n",
        "  y_sum = 0\n",
        "  for y in np.unique(Y):\n",
        "    if np.sum(Y == y) \u003e y_sum:\n",
        "      y_sum = np.sum(Y == y)\n",
        "      y_index = y\n",
        "\n",
        "  # B\n",
        "  # -\n",
        "  # p(y | c)\n",
        "  ix_c = (C == c_index)  # where C = c_index\n",
        "  p_y_c = np.mean(Y[ix_c] == y_index)\n",
        "\n",
        "  # p(y, W | c) = p(W | y , c)p(y | c)\n",
        "  p_y_W_c = np.zeros((k_W,))\n",
        "  ix_c_y = ix_c * (Y == y_index)  # where C = c_index AND Y = y_index\n",
        "  for i in range(k_W):\n",
        "    p_y_W_c[i] = np.mean(W[ix_c_y] == i) * p_y_c\n",
        "\n",
        "  # p(y, X | c) = p(X | y, c)p(y | c)\n",
        "  p_y_X_c = np.zeros((k_X,))\n",
        "  p_X_y_c = np.zeros((k_X,))\n",
        "  # for later: p(Y | X, c)\n",
        "  p_Y_X = np.zeros((k_Y, k_X))\n",
        "  for i in range(k_X):\n",
        "    p_X_y_c[i] = np.mean(X[ix_c_y] == i)\n",
        "    p_y_X_c[i] = p_X_y_c[i] * p_y_c\n",
        "    for j in range(k_Y):\n",
        "      p_Y_X[j, i] = np.mean(Y[(X == i)] == j)\n",
        "\n",
        "  assert np.abs(np.sum(p_X_y_c) - 1) \u003c small\n",
        "\n",
        "  # p(y, X, W | c) = p(W | y, X, c)p(X | y, c)p(y | c)\n",
        "  p_y_X_W_c = np.zeros((k_X, k_W))\n",
        "  # for later: p(y | X, W, C)\n",
        "  p_y_X_W_C = np.zeros((k_X, k_W, k_C))\n",
        "  # for later: p(Y | x, W)\n",
        "  p_Y_X_W = np.zeros((k_Y, k_X, k_W))\n",
        "  for i in range(k_X):\n",
        "    ix_x_c_y = ix_c_y * (X == i)  # where C = c_index AND Y = y_index AND X = i\n",
        "    for j in range(k_W):\n",
        "      p_y_X_W_c[i, j] = np.mean(W[ix_x_c_y] == j) * p_X_y_c[i] * p_y_c\n",
        "      for k in range(k_C):\n",
        "        p_y_X_W_C[i, j,\n",
        "                  k] = np.mean(Y[(X == i) * (W == j) * (C == k)] == y_index)\n",
        "      for k in range(k_Y):\n",
        "        p_Y_X_W[k, i, j] = np.mean(Y[(X == i) * (W == j)] == k)\n",
        "\n",
        "  B = np.zeros((k_X, k_W))\n",
        "  B[0, 0] = p_y_c\n",
        "  B[0, 1:] = p_y_W_c[:-1]\n",
        "  B[1:, 0] = p_y_X_c[:-1]\n",
        "  B[1:, 1:] = p_y_X_W_c[:-1, :-1]\n",
        "\n",
        "  assert np.isnan(B).any() == False\n",
        "\n",
        "  # A\n",
        "  # -\n",
        "  # p(W | c)\n",
        "  p_W_c = np.zeros((k_W,))\n",
        "  for i in range(k_W):\n",
        "    p_W_c[i] = np.mean(W[ix_c] == i)\n",
        "\n",
        "  assert np.abs(np.sum(p_W_c) - 1) \u003c small\n",
        "\n",
        "  # p(X | c)\n",
        "  p_X_c = np.zeros((k_X,))\n",
        "  # for later: p(X), q(X)\n",
        "  p_X = np.zeros((k_X,))\n",
        "  q_X = np.zeros((k_X,))\n",
        "  for i in range(k_X):\n",
        "    p_X[i] = np.mean(X == i)\n",
        "    q_X[i] = np.mean(X_shift == i)\n",
        "    p_X_c[i] = np.mean(X[ix_c] == i)\n",
        "\n",
        "  assert np.abs(np.sum(p_X_c) - 1) \u003c small\n",
        "\n",
        "  # p(X, W | c) = p(W | X, c)p(X | c)\n",
        "  p_X_W_c = np.zeros((k_X, k_W))\n",
        "  # for later: p(W | X, C), p(W | X)\n",
        "  p_W_X_C = np.zeros((k_W, k_X, k_C))\n",
        "  p_W_X = np.zeros((k_W, k_X))\n",
        "  for i in range(k_X):\n",
        "    ix_x_c = ix_c * (X == i)\n",
        "    for j in range(k_W):\n",
        "      p_X_W_c[i, j] = np.mean(W[ix_x_c] == j) * p_X_c[i]\n",
        "      p_W_X[j, i] = np.mean(W[X == i] == j)\n",
        "      for k in range(k_C):\n",
        "        p_W_X_C[j, i, k] = np.mean(W[(C == k) * (X == i)] == j)\n",
        "\n",
        "  assert (np.abs(np.sum(p_W_X_C[:, 0, 0]) - 1) \u003c small)\n",
        "\n",
        "  A = np.zeros((k_X, k_W))\n",
        "  A[0, 0] = 1.0\n",
        "  A[0, 1:] = p_W_c[:-1]\n",
        "  A[1:, 0] = p_X_c[:-1]\n",
        "  A[1:, 1:] = p_X_W_c[:-1, :-1]\n",
        "\n",
        "  assert np.isnan(A).any() == False\n",
        "\n",
        "  # take psuedo-inverse of A\n",
        "  # if k_X \u003e k_W P^+ = (A'A)^{-1}A'\n",
        "  A_inv = np.linalg.pinv(A)\n",
        "  AiB = np.dot(A_inv, B)  # (k_W, k_W)\n",
        "  Delta, H = np.linalg.eig(AiB)\n",
        "\n",
        "  H_inv = np.linalg.pinv(H)  # H_inv is (k_U, k_W)\n",
        "  # get scaling vector e\n",
        "  e = 1 / H_inv[:, 0]\n",
        "  # multiply by H_inv to get S\n",
        "  S = np.dot(np.diag(e), H_inv)\n",
        "  p_W_U = np.column_stack(\n",
        "      (S[:, 1:], 1 - np.sum(S[:, 1:], axis=1))).transpose()  # (k_W, k_U)\n",
        "\n",
        "  # p(U | X) = p(W | U)^{-1}p(W | X)\n",
        "  # -----------------------------------\n",
        "  p_U_X = np.dot(np.linalg.pinv(p_W_U), p_W_X)\n",
        "  assert np.isnan(p_U_X).any() == False\n",
        "\n",
        "  # q(U)/p(U) = p(U | X)^{-1}[q(X)/p(X)]\n",
        "  # ------------------------------------\n",
        "  q_U_p_U = np.dot(q_X / p_X, np.linalg.pinv(p_U_X))\n",
        "  assert np.isnan(q_U_p_U).any() == False\n",
        "  # p(Y | x, U) = p(Y | x, W)p(U | x, W)^{-1}\n",
        "  # p(U | x, W) = (p(W | U) * p(U | x)) / p(W | x)\n",
        "  q_Y_X_ours = np.zeros((k_Y, k_X))\n",
        "  for i in range(k_X):\n",
        "    # p(U | x, W) = (p(W | U) * p(U | x)) / p(W | x)\n",
        "    # ----------------------------------------------\n",
        "    # p(W | U), from above\n",
        "    # p(U | x), from above\n",
        "    # p(W | x), from above\n",
        "\n",
        "    p_U_x_W = ((p_W_U * p_U_X[:, i].reshape((1, k_U))) / p_W_X[:, i].reshape(\n",
        "        (k_W, 1))).transpose()\n",
        "\n",
        "    # p(Y | x, U) = p(Y | x, W)p(U | x, W)^{-1}\n",
        "    # -----------------------------------------\n",
        "    # p(Y | x, W), from above\n",
        "    p_Y_x_U = np.dot(p_Y_X_W[:, i, :], np.linalg.pinv(p_U_x_W))\n",
        "\n",
        "    # q(Y | x) = p(Y | x, U) * p(U | x) * q(U)/p(U) / Z, Z is a normalizing constant to ensure np.dot(1', q(Y | x)) = 1\n",
        "    # -----------------------------------------------------------------------------------------------------------------\n",
        "    unnorm = np.dot(p_Y_x_U, np.dot(np.diag(q_U_p_U,), p_U_X[:, i]))\n",
        "    q_Y_X_ours[:, i] = unnorm / unnorm.sum()  # normalize\n",
        "\n",
        "  # baseline: standard estimator - p_Y_X_c\n",
        "  # oracle: q_Y_X_c\n",
        "  q_Y_X = np.zeros((k_Y, k_X))\n",
        "  q_Y_X_test = np.zeros((k_Y, k_X))\n",
        "\n",
        "  for i in range(k_Y):\n",
        "    for j in range(k_X):\n",
        "      q_Y_X[i, j] = np.mean(Y_shift[(X_shift == j)] == i)\n",
        "      q_Y_X_test[i, j] = np.mean(Y_shift_test[(X_shift_test == j)] == i)\n",
        "\n",
        "  baseline_pYX = p_Y_X\n",
        "  our_qYX = q_Y_X_ours\n",
        "  train_qYX = q_Y_X\n",
        "  test_qYX = q_Y_X_test\n",
        "\n",
        "  return baseline_pYX, our_qYX, train_qYX, test_qYX"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hAx8k09Xc05H"
      },
      "source": [
        "## W=1 (high noise)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WW1ZG_UVbcqt"
      },
      "outputs": [],
      "source": [
        "folder_id = './tmp_data'\n",
        "filename_source = \"synthetic_multivariate_num_samples_10000_w_coeff_1_p_u_0_0.9.csv\"\n",
        "filename_target = \"synthetic_multivariate_num_samples_10000_w_coeff_1_p_u_0_0.1.csv\"\n",
        "data_df_source = pd.read_csv(os.path.join(folder_id, filename_source))\n",
        "data_df_target = pd.read_csv(os.path.join(folder_id, filename_target))\n",
        "data_dict_source = extract_from_df_nested(data_df_source)\n",
        "data_dict_target = extract_from_df_nested(data_df_target)\n",
        "data_dict_all_w1 = dict(source=data_dict_source, target=data_dict_target)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sP3M-V9qZue7"
      },
      "outputs": [],
      "source": [
        "X, C, Y, W, U, X_shift, Y_shift, X_shift_test, Y_shift_test = discretize_data(data_dict_all_w1, is_univariate=False)\n",
        "baseline_pYX, our_qYX, train_qYX, test_qYX = estimate_q_y_x(\n",
        "    X, C, Y, W, U, X_shift, Y_shift, X_shift_test, Y_shift_test)\n",
        "rmse_baseline = rmse(baseline_pYX, test_qYX)\n",
        "rmse_ours = rmse(our_qYX, test_qYX)\n",
        "rmse_sampling = rmse(train_qYX, test_qYX)\n",
        "print(rmse_baseline)\n",
        "print(rmse_ours)\n",
        "print(rmse_sampling)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "507Fg-Ycc4Tp"
      },
      "source": [
        "## W=3 (high noise)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_Ic3WgF-c6gh"
      },
      "outputs": [],
      "source": [
        "folder_id = \"./tmp_data\" # 20221012\n",
        "filename_source = \"synthetic_multivariate_num_samples_10000_w_coeff_3_p_u_0_0.9.csv\"\n",
        "filename_target = \"synthetic_multivariate_num_samples_10000_w_coeff_3_p_u_0_0.1.csv\"\n",
        "data_df_source = pd.read_csv(os.path.join(folder_id, filename_source))\n",
        "data_df_target = pd.read_csv(os.path.join(folder_id, filename_target))\n",
        "data_dict_source = extract_from_df_nested(data_df_source)\n",
        "data_dict_target = extract_from_df_nested(data_df_target)\n",
        "data_dict_all_w3 = dict(source=data_dict_source, target=data_dict_target)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8DW_0nSLc_LY"
      },
      "outputs": [],
      "source": [
        "X, C, Y, W, U, X_shift, Y_shift, X_shift_test, Y_shift_test = discretize_data(data_dict_all_w3, is_univariate=False)\n",
        "baseline_pYX, our_qYX, train_qYX, test_qYX = estimate_q_y_x(\n",
        "    X, C, Y, W, U, X_shift, Y_shift, X_shift_test, Y_shift_test)\n",
        "rmse_baseline = rmse(baseline_pYX, test_qYX)\n",
        "rmse_ours = rmse(our_qYX, test_qYX)\n",
        "rmse_sampling = rmse(train_qYX, test_qYX)\n",
        "print(rmse_baseline)\n",
        "print(rmse_ours)\n",
        "print(rmse_sampling)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.10.9"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
